import tensorflow as tf
from tensorflow.keras import Model
from typing import List, Union
from .traffic_model_trainner import TrafficModelTrainner
from .traffic_model_param import TrafficModelParam


class TrafficModelUser:
    g_model: Union[Model, None]
    d_model: Union[Model, None]
    noise_dim: int
    traffic_dim: int

    def __init__(self, g_path: str, d_path: str, param: TrafficModelParam, trainner: TrafficModelTrainner = None):
        """
        ”模型使用者“ 构造函数
        :param g_path: 生成器模型存储路径
        :param d_path: 判别器模型存储路径
        :param trainner: 训练器
        """
        # 参数的获取
        print("TrafficModelUser Param Initialization.....")
        self.noise_dim = param.noise_dim
        self.traffic_dim = param.traffic_dim
        # 从model_path中读取模型
        try:
            self.g_model = tf.keras.models.load_model(g_path)
        except IOError:
            self.g_model = None
            pass
        try:
            self.d_model = tf.keras.models.load_model(d_path)
        except IOError:
            self.d_model = None
            pass
        # 如果有训练者
        if trainner is not None:
            print("Before Training......")
            # User将模型给予Trainner
            trainner.g_model = self.g_model
            trainner.d_model = self.d_model
            print("now g_model", trainner.g_model)
            print("now d_model", trainner.d_model)
            trainner.train()
            print("Trainning Save.......")
            print("g_model path", g_path)
            print("d_model path", d_path)
            # Trainner将模型给与User
            self.g_model = trainner.g_model
            self.d_model = trainner.d_model
            trainner.g_model.save(g_path)
            trainner.d_model.save(d_path)

    def generate(self, z) -> List[float]:
        generate_traffic_spans = self.g_model(z)
        # generate_traffic_spans = tf.random.normal([1, 10])
        return generate_traffic_spans.numpy().tolist()[0]

    def analysis(self, generate_traffic_spans: List[float]):
        generate_traffic_tensor = tf.constant(generate_traffic_spans)
        dim = generate_traffic_tensor.shape[0]
        generate_traffic_tensor = tf.reshape(generate_traffic_tensor, [1, dim])
        restored_noise = self.d_model(generate_traffic_tensor)
        return restored_noise

    def create_noise(self):
        z = tf.random.uniform([1, self.noise_dim])
        return z
